#!/bin/bash
#
# KWDB SDK Performance Comparison Script
# Auto-validates SDK performance >= 1.5x ODBC/JDBC baseline
#

set -e

# ==============================================================================
# Configuration
# ==============================================================================

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
BUILD_DIR="$PROJECT_ROOT/build"

# Test parameters
TOTAL_RECORDS=100000
BATCH_SIZE=500
CONCURRENCY=4
TABLE_NAME="sdk_bench"

# Minimum performance ratio requirement
MIN_RATIO=1.5

# Output files
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
REPORT_DIR="$PROJECT_ROOT/docs"
REPORT_FILE="$REPORT_DIR/PERF_STRESS_TEST_REPORT_${TIMESTAMP}.md"
TEMP_RESULTS="/tmp/kwdb_bench_results_${TIMESTAMP}.txt"

# Colors
GREEN='\033[0;32m'
RED='\033[0;31m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# ==============================================================================
# Helper Functions
# ==============================================================================

log() {
    echo -e "${GREEN}[$(date +'%Y-%m-%d %H:%M:%S')]${NC} $*"
}

error() {
    echo -e "${RED}[ERROR]${NC} $*" >&2
}

warn() {
    echo -e "${YELLOW}[WARN]${NC} $*"
}

check_command() {
    if ! command -v "$1" &> /dev/null; then
        error "Required command not found: $1"
        exit 1
    fi
}

# Clean database table
clean_table() {
    log "Cleaning table $TABLE_NAME..."

    # Try to connect to KWDB
    if command -v psql &> /dev/null; then
        PGPASSWORD=root psql -U root -h 127.0.0.1 -p 26257 -d kwdb \
            -c "TRUNCATE TABLE $TABLE_NAME;" 2>/dev/null || {
            warn "Failed to truncate via psql, table might not exist"
        }
    else
        warn "psql not found, skipping table cleanup"
    fi
}

# Extract QPS from output
extract_qps() {
    local output="$1"
    local qps=$(echo "$output" | grep -i "QPS:" | awk '{print $2}' | tr -d ',' | head -1)

    if [ -z "$qps" ]; then
        qps=$(echo "$output" | grep -oP "[\d,]+\s+points?/s" | grep -oP "[\d,]+" | tr -d ',' | head -1)
    fi

    echo "$qps"
}

# ==============================================================================
# Benchmark Functions
# ==============================================================================

run_odbc_baseline() {
    log "Running ODBC baseline benchmark..."

    clean_table

    local output=$(python3 "$SCRIPT_DIR/python/odbc_bench.py" \
        --threads $CONCURRENCY \
        --batch $BATCH_SIZE \
        --total $TOTAL_RECORDS \
        2>&1)

    local qps=$(extract_qps "$output")

    if [ -z "$qps" ] || [ "$qps" -eq 0 ]; then
        error "Failed to extract ODBC QPS"
        echo "$output" >&2
        return 1
    fi

    echo "$qps"
}

run_sdk_unary() {
    log "Running SDK Unary mode..."

    clean_table

    local output=$("$BUILD_DIR/bench/cpp/bench_write" \
        $TOTAL_RECORDS $BATCH_SIZE $CONCURRENCY \
        2>&1)

    local qps=$(extract_qps "$output")

    if [ -z "$qps" ]; then
        warn "Failed to extract SDK Unary QPS"
        return 1
    fi

    echo "$qps"
}

run_sdk_stream_tcp() {
    log "Running SDK Stream (TCP) mode..."

    clean_table

    local output=$("$BUILD_DIR/bench/cpp/bench_write_new" \
        --mode stream \
        --transport tcp \
        --host 127.0.0.1 \
        --port 9090 \
        --records $TOTAL_RECORDS \
        --batch-size $BATCH_SIZE \
        --concurrency $CONCURRENCY \
        2>&1)

    local qps=$(extract_qps "$output")

    if [ -z "$qps" ]; then
        warn "Failed to extract SDK Stream TCP QPS"
        return 1
    fi

    echo "$qps"
}

run_sdk_stream_uds() {
    log "Running SDK Stream (UDS) mode..."

    clean_table

    # Check if UDS socket exists
    local uds_path="${KWDB_UDS_PATH:-/tmp/kwdb-ingestd.sock}"
    if [ ! -S "$uds_path" ]; then
        warn "UDS socket not found at $uds_path, skipping"
        echo "0"
        return 0
    fi

    local output=$("$BUILD_DIR/bench/cpp/bench_write_new" \
        --mode stream \
        --transport uds \
        --path "$uds_path" \
        --records $TOTAL_RECORDS \
        --batch-size $BATCH_SIZE \
        --concurrency $CONCURRENCY \
        2>&1)

    local qps=$(extract_qps "$output")

    if [ -z "$qps" ]; then
        warn "Failed to extract SDK Stream UDS QPS"
        return 1
    fi

    echo "$qps"
}

# ==============================================================================
# Main Execution
# ==============================================================================

main() {
    log "KWDB SDK Performance Comparison"
    log "================================"
    log ""
    log "Configuration:"
    log "  Total Records: $TOTAL_RECORDS"
    log "  Batch Size: $BATCH_SIZE"
    log "  Concurrency: $CONCURRENCY"
    log "  Table: $TABLE_NAME"
    log ""

    # Check prerequisites
    check_command python3

    if [ ! -f "$BUILD_DIR/bench/cpp/bench_write_new" ]; then
        error "bench_write_new not found. Please build the project first."
        exit 1
    fi

    # Initialize results file
    > "$TEMP_RESULTS"

    # Run benchmarks
    log "Starting benchmarks..."
    echo ""

    # 1. ODBC Baseline
    ODBC_QPS=$(run_odbc_baseline)
    log "ODBC Baseline QPS: $ODBC_QPS"
    echo "ODBC_BASELINE=$ODBC_QPS" >> "$TEMP_RESULTS"

    sleep 2

    # 2. SDK Unary
    UNARY_QPS=$(run_sdk_unary || echo "0")
    if [ "$UNARY_QPS" != "0" ]; then
        log "SDK Unary QPS: $UNARY_QPS"
        UNARY_RATIO=$(echo "scale=2; $UNARY_QPS / $ODBC_QPS" | bc)
        log "SDK Unary / ODBC Ratio: ${UNARY_RATIO}x"
        echo "SDK_UNARY=$UNARY_QPS" >> "$TEMP_RESULTS"
        echo "UNARY_RATIO=$UNARY_RATIO" >> "$TEMP_RESULTS"
    fi

    sleep 2

    # 3. SDK Stream TCP
    STREAM_TCP_QPS=$(run_sdk_stream_tcp || echo "0")
    if [ "$STREAM_TCP_QPS" != "0" ]; then
        log "SDK Stream TCP QPS: $STREAM_TCP_QPS"
        STREAM_TCP_RATIO=$(echo "scale=2; $STREAM_TCP_QPS / $ODBC_QPS" | bc)
        log "SDK Stream TCP / ODBC Ratio: ${STREAM_TCP_RATIO}x"
        echo "SDK_STREAM_TCP=$STREAM_TCP_QPS" >> "$TEMP_RESULTS"
        echo "STREAM_TCP_RATIO=$STREAM_TCP_RATIO" >> "$TEMP_RESULTS"
    fi

    sleep 2

    # 4. SDK Stream UDS
    STREAM_UDS_QPS=$(run_sdk_stream_uds || echo "0")
    if [ "$STREAM_UDS_QPS" != "0" ]; then
        log "SDK Stream UDS QPS: $STREAM_UDS_QPS"
        STREAM_UDS_RATIO=$(echo "scale=2; $STREAM_UDS_QPS / $ODBC_QPS" | bc)
        log "SDK Stream UDS / ODBC Ratio: ${STREAM_UDS_RATIO}x"
        echo "SDK_STREAM_UDS=$STREAM_UDS_QPS" >> "$TEMP_RESULTS"
        echo "STREAM_UDS_RATIO=$STREAM_UDS_RATIO" >> "$TEMP_RESULTS"
    fi

    echo ""
    log "All benchmarks completed!"

    # Calculate minimum ratio
    MIN_ACTUAL_RATIO=999999
    for ratio_var in UNARY_RATIO STREAM_TCP_RATIO STREAM_UDS_RATIO; do
        ratio_value=$(grep "^${ratio_var}=" "$TEMP_RESULTS" | cut -d'=' -f2)
        if [ -n "$ratio_value" ] && [ "$ratio_value" != "0" ]; then
            # Compare ratios
            if (( $(echo "$ratio_value < $MIN_ACTUAL_RATIO" | bc -l) )); then
                MIN_ACTUAL_RATIO=$ratio_value
            fi
        fi
    done

    echo ""
    log "Results Summary:"
    log "================"
    log "ODBC Baseline: $ODBC_QPS pts/s"
    [ "$UNARY_QPS" != "0" ] && log "SDK Unary: $UNARY_QPS pts/s (${UNARY_RATIO}x)"
    [ "$STREAM_TCP_QPS" != "0" ] && log "SDK Stream TCP: $STREAM_TCP_QPS pts/s (${STREAM_TCP_RATIO}x)"
    [ "$STREAM_UDS_QPS" != "0" ] && log "SDK Stream UDS: $STREAM_UDS_QPS pts/s (${STREAM_UDS_RATIO}x)"
    echo ""
    log "Minimum Ratio: ${MIN_ACTUAL_RATIO}x"

    # Validate minimum ratio
    echo ""
    if (( $(echo "$MIN_ACTUAL_RATIO >= $MIN_RATIO" | bc -l) )); then
        log "${GREEN}✅ PASS: Minimum ratio (${MIN_ACTUAL_RATIO}x) >= ${MIN_RATIO}x${NC}"
        PASS_STATUS="PASS ✅"
        EXIT_CODE=0
    else
        error "${RED}❌ FAIL: Minimum ratio (${MIN_ACTUAL_RATIO}x) < ${MIN_RATIO}x${NC}"
        PASS_STATUS="FAIL ❌"
        EXIT_CODE=1
    fi

    # Generate report
    generate_report

    log "Report saved to: $REPORT_FILE"

    exit $EXIT_CODE
}

generate_report() {
    log "Generating performance report..."

    mkdir -p "$REPORT_DIR"

    cat > "$REPORT_FILE" << EOF
# KWDB SDK Performance Stress Test Report

## Test Information

- **Test Time**: $(date '+%Y-%m-%d %H:%M:%S')
- **Test Environment**: $(uname -s) $(uname -r)
- **Total Records**: $TOTAL_RECORDS
- **Batch Size**: $BATCH_SIZE
- **Concurrency**: $CONCURRENCY
- **Table Name**: $TABLE_NAME

## Performance Results

### Throughput Comparison

| Test Scenario | QPS (points/sec) | vs ODBC Baseline | Status |
|---------------|------------------|------------------|--------|
| **ODBC Baseline** | $ODBC_QPS | 1.00x | Reference |
EOF

    if [ "$UNARY_QPS" != "0" ]; then
        cat >> "$REPORT_FILE" << EOF
| SDK Unary | $UNARY_QPS | **${UNARY_RATIO}x** | $([ $(echo "$UNARY_RATIO >= $MIN_RATIO" | bc -l) -eq 1 ] && echo "✅" || echo "❌") |
EOF
    fi

    if [ "$STREAM_TCP_QPS" != "0" ]; then
        cat >> "$REPORT_FILE" << EOF
| SDK Stream (TCP) | $STREAM_TCP_QPS | **${STREAM_TCP_RATIO}x** | $([ $(echo "$STREAM_TCP_RATIO >= $MIN_RATIO" | bc -l) -eq 1 ] && echo "✅" || echo "❌") |
EOF
    fi

    if [ "$STREAM_UDS_QPS" != "0" ]; then
        cat >> "$REPORT_FILE" << EOF
| SDK Stream (UDS) | $STREAM_UDS_QPS | **${STREAM_UDS_RATIO}x** | $([ $(echo "$STREAM_UDS_RATIO >= $MIN_RATIO" | bc -l) -eq 1 ] && echo "✅" || echo "❌") |
EOF
    fi

    cat >> "$REPORT_FILE" << EOF

### Performance Acceptance

- **Requirement**: SDK Performance ≥ $MIN_RATIO × ODBC Baseline
- **Minimum Ratio**: **${MIN_ACTUAL_RATIO}x**
- **Acceptance Result**: **${PASS_STATUS}**

## Test Configuration

### Server Configuration

\`\`\`bash
# Environment Variables
KWDB_BATCH_SIZE=${KWDB_BATCH_SIZE:-800}
KWDB_FLUSH_INTERVAL_MS=${KWDB_FLUSH_INTERVAL_MS:-20}
WRITER_THREADS=${WRITER_THREADS:-4}
KWDB_UDS_PATH=${KWDB_UDS_PATH:-/tmp/kwdb-ingestd.sock}
\`\`\`

### Table Creation Statement

\`\`\`sql
CREATE TABLE IF NOT EXISTS $TABLE_NAME (
    ts TIMESTAMPTZ NOT NULL,
    host VARCHAR(50),
    region VARCHAR(50),
    usage DOUBLE PRECISION,
    temperature DOUBLE PRECISION,
    PRIMARY KEY (ts)
);
\`\`\`

## Reproduction Commands

### ODBC Baseline
\`\`\`bash
python3 bench/python/odbc_bench.py \\
    --threads $CONCURRENCY \\
    --batch $BATCH_SIZE \\
    --total $TOTAL_RECORDS
\`\`\`

### SDK Unary
\`\`\`bash
./build/bench/cpp/bench_write \\
    $TOTAL_RECORDS $BATCH_SIZE $CONCURRENCY
\`\`\`

### SDK Stream (TCP)
\`\`\`bash
./build/bench/cpp/bench_write_new \\
    --mode stream \\
    --transport tcp \\
    --host 127.0.0.1 --port 9090 \\
    --records $TOTAL_RECORDS \\
    --batch-size $BATCH_SIZE \\
    --concurrency $CONCURRENCY
\`\`\`

### SDK Stream (UDS)
\`\`\`bash
./build/bench/cpp/bench_write_new \\
    --mode stream \\
    --transport uds \\
    --path /tmp/kwdb-ingestd.sock \\
    --records $TOTAL_RECORDS \\
    --batch-size $BATCH_SIZE \\
    --concurrency $CONCURRENCY
\`\`\`

## Conclusion

EOF

    if [ "$EXIT_CODE" -eq 0 ]; then
        cat >> "$REPORT_FILE" << EOF
✅ **KWDB SDK Performance Test PASSED**

All test scenarios have met or exceeded the ${MIN_RATIO}× ODBC baseline requirement. Minimum ratio achieved: **${MIN_ACTUAL_RATIO}x**.

EOF
    else
        cat >> "$REPORT_FILE" << EOF
❌ **KWDB SDK Performance Test FAILED**

Some test scenarios did not meet the ${MIN_RATIO}× ODBC baseline requirement. Minimum ratio achieved: **${MIN_ACTUAL_RATIO}x**.

Performance optimization or test configuration review required.
EOF
    fi

    cat >> "$REPORT_FILE" << EOF

---

*Report generated at: $(date '+%Y-%m-%d %H:%M:%S')*
EOF
}

# Run main function
main "$@"
